package com.amumu.drama.common.config.redis;

import com.google.common.base.Charsets;
import com.google.common.hash.Hashing;

/**
 * 高并发请求限制器
 * 使用场景：某资源有唯一ID在某段时间内并发请求最多100个
 *
 * 原理：使用redis的sorted set记录资源唯一ID个数
 *
 * see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#concurrent-requests-limiter
 * Created by freethinker on 2020-07-01.
 */
public class JedisConcurrentRequestLimiter {

    // @formatter:off
    private static final String SCRIPT=
        "local key = KEYS[1] " +
            " " +
            "local capacity = tonumber(ARGV[1]) " +          // 容量
            "local timestamp = tonumber(ARGV[2]) " +         // 时间戳
            "local id = ARGV[3] " +                          // 唯一ID
            "local ttl = tonumber(ARGV[4]) " +         // 时间戳
            "redis.call('zremrangebyscore', key, '-inf',timestamp-ttl) " +  // 清除过期id
            " " +
            "local count = redis.call('zcard', key) " +       // 获取当前数量
            "local allowed = count < capacity " +             // 是否允许请求
            " " +
            "if allowed then " +
            "  redis.call('zadd', key, timestamp, id) " +     // 增加到桶
            "  redis.call('expire', key, ttl*2) " +             // 设置过期时间
            "end " +
            " " +
            "return allowed";                                 // 相当于原子性的zcard+zadd
    // @formatter:on

    /** 限流字段名 */
    private String limitRateName;

    private JedisTemplate jedisTemplate;

    /** 时间段，单位：秒 */
    private int ttl;

    /** 桶的容量，即个数 */
    private int burstCapacity;

    JedisConcurrentRequestLimiter(JedisRateLimiterConfig config, JedisTemplate jedisTemplate) {
        this.limitRateName =  Hashing.md5().hashString(config.getLimitRateName(), Charsets.UTF_8).toString().substring(8, 16);
        this.jedisTemplate = jedisTemplate;
        this.ttl = config.getLimitSpanSeconds();
        this.burstCapacity = config.getLimitCount();
    }

    public boolean tryAcquire(String uniqueId) {
        if (uniqueId == null || uniqueId.isEmpty()) {
            return false;
        }
        long timestamp = System.currentTimeMillis() / 1000;

        String reidsKey = getRedisKey();

        return jedisTemplate.execute(jedis -> null != jedis.eval(SCRIPT, 1, reidsKey, "" + burstCapacity, "" + timestamp, uniqueId, "" + ttl));
    }

    /**
     * 获取当前数量
     */
    public Long getCurrentCount() {

        long timestamp = System.currentTimeMillis() / 1000;

        String reidsKey = getRedisKey();
        // 清除过期的id
        jedisTemplate.sortedSet().zremrangeByScore(reidsKey, "-inf", "" + (timestamp - ttl));

        return jedisTemplate.sortedSet().zcard(reidsKey);
    }

    /**
     * 清除桶
     */
    public void clear() {
        jedisTemplate.key().del(getRedisKey());
    }

    private String getRedisKey() {
        return "cl_" + limitRateName;
    }
}
